from torch.utils.data import Dataset
import numpy as np
from ..dataset import build_preprocess_op


class BaseDataset(Dataset):

    def __init__(self, path, preprocess_pipe, mode='train'):
        self.data_path = path
        self.input_data = []
        self.annotation = []
        self.train_indexes = []
        self.test_indexes = []
        self.cur_indexes = None
        self.mode = None
        self.data_pipe = None
        self.set_mode(mode)
        self.load_data()
        self.preprocess(preprocess_pipe)
        print('---loaded data---')

    def load_data(self):
        raise Exception('BaseDataset does not implement load_data function')

    def __len__(self):
        return len(self.cur_indexes)

    def __getitem__(self, idx):
        img = self.input_data[self.cur_indexes[idx]]
        annotation = self.annotation[self.cur_indexes[idx]]
        if img.dtype != 'float32':
            img = img.astype(np.float32)
        for preprocess_op in self.data_pipe:
            img = preprocess_op.process(img)
        img = img.transpose(2, 0, 1)
        return img, annotation

    def preprocess(self, preprocess_pipe_cfgs):
        preprocess_pipe = []
        for data_preprocess_cfg in preprocess_pipe_cfgs:
            preprocess_pipe.append(build_preprocess_op(data_preprocess_cfg))
        for i, img in enumerate(self.input_data):
            if img.dtype != 'float32':
                img = img.astype(np.float32)
            for preprocess_op in preprocess_pipe:
                img = preprocess_op.process(img)
            self.input_data[i] = img

    def set_data_pipe(self, data_pipe):
        self.data_pipe = data_pipe

    def set_mode(self, mode):
        self.mode = mode
        if self.mode == 'train':
            self.cur_indexes = self.train_indexes
        else:
            self.cur_indexes = self.test_indexes
